#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import json
import chardet
import re
import time
import base64
from multiprocessing import Pool, cpu_count
from collections import defaultdict

# 二进制文件的扩展名列表（常见的二进制文件类型）
BINARY_EXTENSIONS = {
    # 图像文件
    '.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.ico', '.webp',
    # 音频文件
    '.mp3', '.wav', '.ogg', '.flac', '.aac', '.wma',
    # 视频文件
    '.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm',
    # 压缩文件
    '.zip', '.rar', '.7z', '.tar', '.gz', '.bz2', '.xz',
    # 可执行文件
    '.exe', '.dll', '.so', '.dylib', '.bin',
    # 文档和其他二进制格式
    '.pdf', '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx',
    '.class', '.pyc', '.o', '.a', '.lib'
}

def is_likely_binary(file_path, content_sample=None):
    """
    判断文件是否可能是二进制文件
    
    参数:
    - file_path: 文件路径
    - content_sample: 文件内容的样本，如果已经读取则可以提供以避免重复读取
    
    返回:
    - 如果可能是二进制文件则返回True，否则返回False
    """
    # 根据扩展名快速判断
    _, ext = os.path.splitext(file_path)
    if ext.lower() in BINARY_EXTENSIONS:
        return True
    
    # 如果没有提供内容样本，读取文件开头部分
    if content_sample is None:
        try:
            with open(file_path, 'rb') as f:
                content_sample = f.read(1024)  # 读取文件前1KB内容
        except Exception:
            return False
    
    # 检查内容中是否包含大量的空字节或非ASCII字符
    # 这是识别二进制文件的简单方法
    try:
        null_bytes = content_sample.count(b'\x00')
        if null_bytes > len(content_sample) * 0.1:  # 如果空字节超过10%
            return True
        
        # 尝试解码为UTF-8，如果失败则可能是二进制文件
        content_sample.decode('utf-8')
        return False
    except UnicodeDecodeError:
        # 尝试检测编码
        encoding = chardet.detect(content_sample)
        if encoding['confidence'] < 0.7:  # 如果检测置信度低，可能是二进制文件
            return True
        return False

def detect_encoding(file_path):
    """
    检测文件的编码方式
    """
    try:
        with open(file_path, 'rb') as f:
            content_sample = f.read(4096)  # 读取更多内容以提高检测精度
            
            # 如果看起来是二进制文件，返回None
            if is_likely_binary(file_path, content_sample):
                return None
                
            result = chardet.detect(content_sample)
        return result['encoding']
    except Exception:
        return None

def read_file_content(file_path, include_binary=False, binary_encode_base64=True):
    """
    读取文件内容，处理编码问题
    
    参数:
    - file_path: 文件路径
    - include_binary: 是否包含二进制文件
    - binary_encode_base64: 是否对二进制内容进行base64编码
    
    返回:
    - 包含文件内容和元数据的字典
    """
    try:
        # 读取小样本来检测文件类型
        with open(file_path, 'rb') as f:
            content_sample = f.read(1024)
        
        is_binary = is_likely_binary(file_path, content_sample)
        
        if is_binary:
            if not include_binary:
                return {
                    'success': True,
                    'is_binary': True,
                    'content': None,
                    'message': '二进制文件，内容未包含'
                }
            else:
                # 处理二进制文件
                with open(file_path, 'rb') as f:
                    binary_content = f.read()
                
                if binary_encode_base64:
                    # 将二进制内容编码为base64字符串
                    encoded_content = base64.b64encode(binary_content).decode('ascii')
                    return {
                        'success': True,
                        'is_binary': True,
                        'encoding': 'base64',
                        'content': encoded_content,
                        'original_size': len(binary_content)
                    }
                else:
                    # 直接包含二进制内容（不推荐，可能会导致JSON序列化问题）
                    return {
                        'success': True,
                        'is_binary': True,
                        'content': binary_content,
                        'warning': '直接包含二进制内容可能导致JSON序列化问题'
                    }
        else:
            # 处理文本文件
            encoding = detect_encoding(file_path)
            
            # 如果检测不到编码或编码不支持，使用UTF-8编码
            if not encoding or encoding.lower() == 'ascii':
                encoding = 'utf-8'
            
            with open(file_path, 'r', encoding=encoding, errors='replace') as f:
                content = f.read()
            return {'success': True, 'is_binary': False, 'content': content, 'encoding': encoding}
    except Exception as e:
        # 如果读取失败，返回错误信息
        return {'success': False, 'error': str(e), 'path': file_path}

def get_all_files(directory, exclude_dirs=None, include_extensions=None, exclude_extensions=None):
    """
    获取指定目录下的所有文件路径
    
    参数:
    - directory: 要搜索的目录
    - exclude_dirs: 要排除的目录列表
    - include_extensions: 要包含的文件扩展名列表，如果为None，则包含所有文件
    - exclude_extensions: 要排除的文件扩展名列表
    
    返回:
    - 文件路径列表
    """
    if exclude_dirs is None:
        exclude_dirs = ['.git', '__pycache__', 'node_modules']
    
    all_files = []
    
    for root, dirs, files in os.walk(directory):
        # 排除指定目录
        dirs[:] = [d for d in dirs if d not in exclude_dirs]
        
        for file in files:
            file_path = os.path.join(root, file)
            
            # 获取文件扩展名
            _, ext = os.path.splitext(file)
            ext = ext.lower()
            
            # 如果有指定排除的扩展名，则排除这些扩展名的文件
            if exclude_extensions is not None and ext in exclude_extensions:
                continue
            
            # 如果有指定包含的扩展名，则只处理这些扩展名的文件
            if include_extensions is not None and ext not in include_extensions:
                continue
            
            all_files.append(file_path)
    
    return all_files

def process_file(args):
    """
    处理单个文件，用于多进程并行处理
    """
    file_path, max_file_size, pattern, include_binary, binary_encode_base64 = args
    
    # 检查文件大小
    try:
        if os.path.getsize(file_path) > max_file_size:
            return file_path, {
                'success': False,
                'error': f'文件大小超过限制 ({max_file_size} 字节)',
                'path': file_path
            }
    except Exception as e:
        return file_path, {'success': False, 'error': str(e), 'path': file_path}
    
    # 读取文件内容
    result = read_file_content(file_path, include_binary, binary_encode_base64)
    
    # 如果是二进制文件且不需要匹配模式，或者没有指定模式，则直接返回结果
    if result.get('is_binary', False) or not pattern or not result.get('success', False):
        return file_path, result
    
    # 如果需要过滤内容
    try:
        matches = re.findall(pattern, result['content'], re.MULTILINE)
        if not matches:
            return file_path, None  # 没有匹配的内容，返回None
        result['matches'] = matches
        result['match_count'] = len(matches)
    except Exception as e:
        result['match_error'] = str(e)
    
    return file_path, result

def get_directory_content(directory, exclude_dirs=None, include_extensions=None, exclude_extensions=None,
                          max_file_size=1024*1024, pattern=None, parallel=True, max_workers=None,
                          include_binary=False, binary_encode_base64=True):
    """
    获取指定目录下的所有文件内容
    
    参数:
    - directory: 要搜索的目录
    - exclude_dirs: 要排除的目录列表
    - include_extensions: 要包含的文件扩展名列表，如果为None，则包含所有文件
    - exclude_extensions: 要排除的文件扩展名列表
    - max_file_size: 最大文件大小，超过此大小的文件将被跳过，默认为1MB
    - pattern: 正则表达式模式，用于过滤文件内容
    - parallel: 是否并行处理文件
    - max_workers: 最大工作进程数，默认为CPU核心数
    - include_binary: 是否包含二进制文件内容
    - binary_encode_base64: 是否对二进制内容进行base64编码
    
    返回:
    - 包含文件内容的字典
    """
    start_time = time.time()
    all_files = get_all_files(directory, exclude_dirs, include_extensions, exclude_extensions)
    result = {}
    
    if parallel and len(all_files) > 10:  # 只有当文件数量足够多时才使用并行处理
        if not max_workers:
            max_workers = max(1, cpu_count() - 1)  # 使用CPU核心数减1，避免占用所有资源
        
        # 准备参数
        args_list = [(f, max_file_size, pattern, include_binary, binary_encode_base64) for f in all_files]
        
        # 并行处理文件
        with Pool(max_workers) as pool:
            for file_path, file_result in pool.imap_unordered(process_file, args_list):
                if file_result is not None:  # 只有当结果不为None时才添加到结果集
                    result[file_path] = file_result
    else:
        # 串行处理文件
        for file_path in all_files:
            _, file_result = process_file((file_path, max_file_size, pattern, include_binary, binary_encode_base64))
            if file_result is not None:
                result[file_path] = file_result
    
    # 添加处理时间信息
    result['__meta__'] = {
        'total_files': len(all_files),
        'processed_files': len(result) - 1 if '__meta__' in result else len(result),
        'processing_time': time.time() - start_time
    }
    
    return result

def generate_summary(result):
    """
    根据获取的文件内容生成摘要信息
    
    参数:
    - result: get_directory_content函数返回的结果
    
    返回:
    - 摘要信息字典
    """
    summary = {
        'total_files': result['__meta__']['processed_files'] if '__meta__' in result else len(result),
        'success_count': 0,
        'error_count': 0,
        'binary_count': 0,
        'text_count': 0,
        'extension_stats': defaultdict(int),
        'encoding_stats': defaultdict(int),
        'total_size': 0,
        'average_size': 0,
        'largest_files': []
    }
    
    file_sizes = []
    
    for file_path, data in result.items():
        if file_path == '__meta__':
            continue
        
        # 统计成功和失败的文件数
        if data.get('success', False):
            summary['success_count'] += 1
            
            # 区分二进制和文本文件
            if data.get('is_binary', False):
                summary['binary_count'] += 1
            else:
                summary['text_count'] += 1
            
            # 统计文件编码
            encoding = data.get('encoding', 'unknown')
            summary['encoding_stats'][encoding] += 1
            
            # 获取文件大小
            try:
                file_size = os.path.getsize(file_path)
                summary['total_size'] += file_size
                file_sizes.append((file_path, file_size))
            except:
                pass
        else:
            summary['error_count'] += 1
        
        # 统计文件扩展名
        _, ext = os.path.splitext(file_path)
        ext = ext.lower() if ext else '(no extension)'
        summary['extension_stats'][ext] += 1
    
    # 计算平均文件大小
    if summary['success_count'] > 0:
        summary['average_size'] = summary['total_size'] / summary['success_count']
    
    # 获取最大的10个文件
    file_sizes.sort(key=lambda x: x[1], reverse=True)
    summary['largest_files'] = [{'path': p, 'size': s} for p, s in file_sizes[:10]]
    
    # 将defaultdict转换为普通dict以便JSON序列化
    summary['extension_stats'] = dict(summary['extension_stats'])
    summary['encoding_stats'] = dict(summary['encoding_stats'])
    
    return summary

def main():
    """
    主函数，处理命令行参数并执行相应操作
    """
    import argparse
    
    parser = argparse.ArgumentParser(description='获取目录下的所有文件内容')
    parser.add_argument('-d', '--directory', default='.', help='要搜索的目录，默认为当前目录')
    parser.add_argument('-o', '--output', help='输出结果到文件')
    parser.add_argument('-e', '--extensions', help='要包含的文件扩展名，用逗号分隔，如 .py,.txt')
    parser.add_argument('--exclude-extensions', help='要排除的文件扩展名，用逗号分隔，如 .jpg,.pdf')
    parser.add_argument('--exclude-dirs', help='要排除的目录，用逗号分隔，如 .git,node_modules')
    parser.add_argument('--max-size', type=int, default=1024*1024, help='最大文件大小，单位为字节，默认为1MB')
    parser.add_argument('-p', '--pattern', help='使用正则表达式过滤文件内容')
    parser.add_argument('--no-parallel', action='store_true', help='禁用并行处理')
    parser.add_argument('--workers', type=int, help='并行处理的工作进程数，默认为CPU核心数')
    parser.add_argument('--summary-only', action='store_true', help='只输出摘要信息，不输出文件内容')
    parser.add_argument('--include-binary', action='store_true', help='包含二进制文件内容')
    parser.add_argument('--no-base64', action='store_true', help='不对二进制内容进行base64编码（不推荐）')
    
    args = parser.parse_args()
    
    # 处理排除目录
    exclude_dirs = None
    if args.exclude_dirs:
        exclude_dirs = args.exclude_dirs.split(',')
    
    # 处理包含扩展名
    include_extensions = None
    if args.extensions:
        include_extensions = [ext if ext.startswith('.') else '.' + ext for ext in args.extensions.split(',')]
    
    # 处理排除扩展名
    exclude_extensions = None
    if args.exclude_extensions:
        exclude_extensions = [ext if ext.startswith('.') else '.' + ext for ext in args.exclude_extensions.split(',')]
    
    # 获取文件内容
    result = get_directory_content(
        directory=args.directory, 
        exclude_dirs=exclude_dirs, 
        include_extensions=include_extensions,
        exclude_extensions=exclude_extensions,
        max_file_size=args.max_size,
        pattern=args.pattern,
        parallel=not args.no_parallel,
        max_workers=args.workers,
        include_binary=args.include_binary,
        binary_encode_base64=not args.no_base64
    )
    
    # 生成摘要
    summary = generate_summary(result)
    
    # 如果只输出摘要
    if args.summary_only:
        output_content = summary
    else:
        # 将摘要添加到结果中
        result['__summary__'] = summary
        output_content = result
    
    # 处理输出
    if args.output:
        with open(args.output, 'w', encoding='utf-8') as f:
            json.dump(output_content, f, ensure_ascii=False, indent=2)
        print(f"结果已保存到 {args.output}")
    else:
        print(json.dumps(output_content, ensure_ascii=False, indent=2))

if __name__ == "__main__":
    main() 